# LangGraph 中通过检查点为任何 StateGraph 提供内存。在创建任何 LangGraph 工作流时，您可以通过以下方式设置它们以持久保存其状态。
import os
from time import sleep
from typing import TypedDict, Annotated, Literal

from dotenv import load_dotenv
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

# 加载环境变量（通常包含 API 密钥等配置信息）
load_dotenv()


# 定义一个状态类，使用 TypedDict 来类型化状态字典
# 包含一个消息列表，使用 Annotated 和 add_messages 来指定消息的合并方式
class MyState(TypedDict):
    messages: Annotated[list, add_messages]


# 定义一个名为 search 的工具函数，用于模拟网络搜索
@tool
def search(query: str):
    """Call to surf the web."""  # 工具描述，用于让 LLM 理解工具用途
    # 这是实际实现的占位符，返回模拟的搜索结果
    return ["The answer to your question lies within"]


# 将工具函数放入列表，方便后续绑定到模型和工具节点
tools = [search]

# 创建一个 ToolNode 实例，传入工具列表，用于执行工具调用
tool_node = ToolNode(tools)

# 创建大语言模型实例，配置相关参数
llm = ChatOpenAI(
    api_key=os.getenv("MODELSCOPE_API_KEY"),  # 从环境变量获取 API 密钥
    openai_api_base=os.getenv("MODELSCOPE_API_BASE"),  # API 基础地址
    model="Qwen/Qwen2.5-VL-72B-Instruct",  # 使用的模型名称
    streaming=True,  # 启用流式输出
    temperature=0.7  # 控制生成随机性的温度参数
)

# 将工具绑定到模型上，使模型能够识别和调用这些工具
bound_model = llm.bind_tools(tools)


# 定义一个路由函数，根据当前状态决定下一个执行的节点
def should_continue(state: MyState) -> Literal["action", "__end__"]:
    """返回下一个要执行的节点"""
    # 获取最后一条消息
    last_message = state["messages"][-1]

    # 检查最后一条消息是否包含工具调用
    # 如果没有函数调用，则结束工作流
    if not last_message.tool_calls:
        return "__end__"

    # 如果有工具调用，则继续执行 action 节点
    return "action"


# 定义调用大语言模型的函数
def call_llm(state: MyState):
    """调用 LLM 并返回响应"""
    # 使用 LLM 处理当前消息列表
    response = llm.invoke(state["messages"])
    # 返回包含响应消息的字典，格式为列表以便添加到现有消息列表中
    return {"messages": response}


# ========== 构建工作流图 ==========

# 创建状态图实例，指定状态类型为 MyState
workflow = StateGraph(MyState)

# 添加两个节点：
# - "agent": 负责调用 LLM 进行推理
# - "action": 负责执行工具调用
workflow.add_node("agent", call_llm)
workflow.add_node("action", tool_node)

# 设置入口节点，从 START 连接到 agent 节点
workflow.add_edge(START, "agent")

# 添加条件边：根据 should_continue 函数的返回值决定从 agent 节点出发的路径
workflow.add_conditional_edges("agent", should_continue)

# 添加普通边：从 action 节点返回到 agent 节点，形成循环
workflow.add_edge("action", "agent")

# 创建内存检查点保存器，用于持久化保存对话状态
memory = MemorySaver()

# 编译工作流，传入检查点保存器以启用状态持久化
app = workflow.compile(checkpointer=memory)

# ========== 执行工作流 ==========

# 设置配置参数，thread_id 用于区分不同的对话线程
config = {"configurable": {"thread_id": "2"}}

# 创建第一条人类消息
input_message = HumanMessage(content="你好，我是继风")

# 在流模式下运行应用程序，实时处理每个事件
print("=== 第一次对话 ===")
for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    # 提取并美化打印每个事件的最后一条消息
    event["messages"][-1].pretty_print()

# 等待2秒，模拟用户思考时间
sleep(2)

# 创建第二条消息，测试记忆功能
print("\n=== 第二次对话（测试记忆） ===")
input_message = HumanMessage(content="我的名字是什么")
for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()

# 等待5秒
sleep(5)

# 模拟切换到新线程（thread_id=3），此时会丢失之前的对话记忆
print("\n=== 第三次对话（新线程，无记忆） ===")
input_message = HumanMessage(content="我叫什么名字?")
for event in app.stream({"messages": [input_message]}, {"configurable": {"thread_id": "3"}}, stream_mode="values"):
    event["messages"][-1].pretty_print()

# 等待5秒
sleep(5)

# 重新回到原来的线程（thread_id=2），此时应该恢复之前的对话记忆
print("\n=== 第四次对话（回到原线程，恢复记忆） ===")
input_message = HumanMessage(content="你忘了吗?")
for event in app.stream({"messages": [input_message]}, {"configurable": {"thread_id": "2"}}, stream_mode="values"):
    event["messages"][-1].pretty_print()